#from _typeshed import Self
#from _typeshed import Self

import gym
import math
from gym.envs.classic_control import rendering
import numpy as np
from gym import spaces
from gym.utils import seeding
import math
import random

def cnt():
    x=random.random()
    if x>0.25 :
        return 1
    else:
        return 0

class GridEnv(gym.Env):
    metadata = {
        'render.modes': ['human', 'rgb_array'], #选择渲染模式
        'video.frames_per_second': 50
    }

    def __init__(self):
        self.n=500                                      #边数
        self.r=250
        self.sheep_V=1*math.pi*self.r/self.n
        self.sheep_v_=[0,0]
        self.sheep_x=0
        self.sheep_y=0

        self.suc=0
        
        self.p=[]                                            ##初始化狼的运动轨迹点
        for i in range(self.n):                                   
             ang = 2*math.pi*i /(self.n) 
             self.p.append([math.cos(ang)*250, math.sin(ang)*250])

        self.dog_fpostion=(250,0)
        self.f=0
        #self.low=np.array[-300,-300]
        #self.high=np.array[300,300]
        
        self.low=np.array([-300,-300,0],dtype=np.float32)
        self.high=np.array([300,300,2*math.pi],dtype=np.float32)
                                        #视图开关工具,为false时视图打开
        self.viewer = False                           #一开始不展示视图
        self.action_space = spaces.Box(low=-0.5*np.pi, high=2.5*np.pi, shape=(1,1), dtype=np.float32) 
        self.observation_space =spaces. Box(self.low, self.high,  dtype=np.float32)  
        #action(动作空间)为1X1数组，范围0-2 pi
        #state(状态空间2)2x2数组,范围都是-50-50
        #self.observation_space=[0,0,0]
        self.dog_dirctcon=1                          #方向控制开关，为1时候狼走逆时针,2时候顺时针，0为停止


    def step(self, action):
        self.dog_fpostion=self.p[self.f]
        jiao_wolf=he(math.atan2(self.dog_fpostion[1],self.dog_fpostion[0]))        #计算狼的位置角度
        jiao_sheep=he(math.atan2(self.sheep_y,self.sheep_x))                   #计算羊的位置角度
        #jiao_wolf=self.he(self.dog_fpostion[0],self.dog_fpostion[1])
        #jiao_sheep=self.he(self.sheep_x,self.sheep_y)    
        res_error=jiao_wolf-jiao_sheep                                             
        
    ################################################################狼的运动：
        #如果角度差在0-180 或者 -360到-180之间，则狼顺时针运动
        if((res_error>0 and res_error<=math.pi) or (res_error>-2*math.pi) and (res_error<=-math.pi)): 
            self.dog_dirctcon=2
    
        elif(res_error==0 or res_error==-2*math.pi):   #如果角度差为0或者360，则不动
            self.dog_dirctcon=0

        #如果角度差在0-180 或者 -360到-180之间，则顺时针运动
        else: 
            self.dog_dirctcon=1
 
        if(self.dog_dirctcon==1):                           #逆时针圆周运动
            if(self.f<(self.n-1)):
                
                self.f=self.f+1              #self.f为计数器
            else:
                self.f=0
        elif(self.dog_dirctcon==2):
            if(self.f>=0):                                  #顺时针圆周运动
                #self.dog_fpostion=self.p[self.f]
                self.f=self.f-1
            else:
                 self.f=self.n-1
        else:self.f=self.f
        


        #############################################################羊的运动代码：
        #a,b=0,2*math.pi
        a=jiao_sheep-0.5*math.pi
        b=jiao_sheep+0.5*math.pi         #动作空间范围
        self.action=a + (b-a)*random.random() #a-------      在动作空间区间内随机数（动作）  
        self.sheep_v_=[math.cos(self.action)*self.sheep_V,math.sin(self.action)*self.sheep_V]     #生成羊的速度矢量
        self.sheep_x+=self.sheep_v_[0]                                                                               
        self.sheep_y+=self.sheep_v_[1]             #生成羊的下一秒位置

       # biao_act=biao(self.action)
        jiao_sheep2=he(math.atan2(self.sheep_y,self.sheep_x))     #羊下一秒位置的角度
        next_state=np.array([self.sheep_x,self.sheep_y,jiao_sheep2])        #生成[x,y,角度]，也就是羊的状态数组
        
        
        sheep_o=math.sqrt(self.sheep_x** 2+self.sheep_y**2)                        #羊和原点的距离
        sheep_dog=math.sqrt((self.sheep_x-self.dog_fpostion[0]) ** 2+(self.sheep_y-self.dog_fpostion[1])**2)     #羊和狼的距离

        
        done = bool(
            sheep_o>=256 
            or sheep_dog<=10
        )

        reward=1*sheep_o+1.5* sheep_dog+ 10 * (self.action)        #设定奖励
        return next_state,reward,done,{}
       # return jiao_sheep,[a,b,self.action],done,{}

    def reset(self):
        self.sheep_v_=[0,0]
        a,b=0,2*math.pi
      
        rnd=a + (b-a)*random.random()
        self.sheep_x=10*math.cos(rnd)
        self.sheep_y=10*math.sin(rnd)
        self.dog_dirctcon=[250,0]
        self.f=0
        self.viewer = None   

        tate=np.array([self.sheep_x,self.sheep_y,0])   
        #return np.array([tate])
        return tate
    
            
    def render(self, mode='human'):
                                #需要修改的代码
        
       
        if self.viewer is None:                   
            self.viewer = rendering.Viewer(600, 600)
            self.viewer.isopen=True
            self.viewer.transform = rendering.Transform(translation=(300,300))         #视图+坐标迁移
            

                                                                         #画一个直径为 500 的基本草地:
            circle = rendering.make_circle(250,filled=False,res=50)

        # 添加一个平移操作:
            circle_transform = rendering.Transform(translation=(0, 0))
            circle.set_color(0, 255, 0)

        # 让草地添加平移这个属性
            circle.add_attr(circle_transform)
            circle.set_linewidth(4)
            self.viewer.add_geom(circle)

            
            self.dog=rendering.make_circle(6,res=50)     #添加狗子！！
            self.dog_transfrom = rendering.Transform(translation=(250,0))
            self.dog.set_color(255, 0, 0)
            self.dog.add_attr(self.dog_transfrom)
            self.viewer.add_geom(self.dog)

            self.sheep=rendering.make_circle(6,res=50)    #添加山羊！！！
            self.sheep_transfrom = rendering.Transform(translation=(self.sheep_x,self.sheep_y))
            self.sheep.set_color(0,0,255)
            self.sheep.add_attr(self.sheep_transfrom)
            self.viewer.add_geom(self.sheep)
        
        self.dog_transfrom.set_translation(self.dog_fpostion[0],self.dog_fpostion[1])
        self.sheep_transfrom.set_translation(self.sheep_x,self.sheep_y)   
     
            
           

            
           
        return self.viewer.render(return_rgb_array=mode == 'rgb_array')


               
    def seed(self, seed=None):
        self.np_random, seed = seeding.np_random(seed)
        return [seed]

    
    
def he(x):
    if x<=0:x+=2*math.pi
    return x

def biao(x):
    if x<0:
        x+=2*math.pi
    elif x>2*math.pi:
        x-=2*math.pi
   